#include "Basis.h"
#include "Hook.h"
#include "Memory.h"
#include "Cpu.h"

ULONG						g_HookReturnAddress;

LIST_ENTRY					g_HookMonitorList;
KSPIN_LOCK					g_HookMonitorListSpinLock;

LIST_ENTRY					g_HookMonitorInfoList;
KSPIN_LOCK					g_HookMonitorInfoListSpinLock;
NPAGED_LOOKASIDE_LIST		g_HookMonitorInfoListLookaside;
HANDLE						g_HookMonitorInfoListEventHandler;
PKEVENT						g_HookMonitorInfoListEvent;

VOID InitializeHook()
{
	InitializeListHead(&g_HookMonitorList);
	KeInitializeSpinLock(&g_HookMonitorListSpinLock);

	InitializeListHead(&g_HookMonitorInfoList);
	KeInitializeSpinLock(&g_HookMonitorInfoListSpinLock);
	ExInitializeNPagedLookasideList(&g_HookMonitorInfoListLookaside, NULL, NULL, 0, sizeof(MONITOR_INFO_ITEM), HOOK_ALLOC_TAG, 0);

	UNICODE_STRING monitorInfoEventName;
	RtlInitUnicodeString(&monitorInfoEventName, MONITOR_INFO_EVENT_NAME);
	g_HookMonitorInfoListEvent = IoCreateNotificationEvent(&monitorInfoEventName, &g_HookMonitorInfoListEventHandler);
	KeClearEvent(g_HookMonitorInfoListEvent);
}

VOID DisposeHook()
{
	EmptyHookLists();
	ExDeleteNPagedLookasideList(&g_HookMonitorInfoListLookaside);
	ZwClose(g_HookMonitorInfoListEventHandler);
}

BOOLEAN ProcessIoCtlForHook(ULONG ioControlCode, PDEVICE_OBJECT deviceObject, PIRP irp, PIO_STACK_LOCATION irpStack)
{
	UNREFERENCED_PARAMETER(deviceObject);

	switch (ioControlCode)
	{
	case IOCTL_HOOK_INSTALL:
		ProcessHook(irp, TRUE);
		break;
	case IOCTL_HOOK_UNINSTALL:
		ProcessHook(irp, FALSE);
		break;
	case IOCTL_HOOK_SET_MONITOR_ITEM:
		ProcessSetMonitorItem(irp, irpStack);
		break;
	case IOCTL_HOOK_GET_MONITOR_INFO:
		ProcessGetMonitorInfo(irp, irpStack);
		break;
	default:
		return FALSE;
	}
	return TRUE;
}

VOID ProcessHook(PIRP irp, BOOLEAN isInstall)
{
	if (isInstall)
	{
		EmptyHookLists();
	}

	// make target address writable
	ULONG addressOfHook = *(PULONG)irp->AssociatedIrp.SystemBuffer;
	ULONG addressOfWrite;
	PMDL pMdl = AllocateWritableMdl(addressOfHook, 16, &addressOfWrite);
	KdPrint(("hook address 0x%x, writable mapped address 0x%x", addressOfHook, addressOfWrite));
	if (!pMdl)
	{
		irp->IoStatus.Status = STATUS_UNSUCCESSFUL;
		return;
	}
	
	// busy cpu and do hooking
	PVOID params[3] = { 0 };
	params[0] = (PVOID)addressOfHook;
	params[1] = (PVOID)addressOfWrite;
	params[2] = (PVOID)isInstall;
	BOOLEAN retVal = BusyCpu(HookInternal, params);
	FreeMdl(pMdl);
	irp->IoStatus.Status = retVal ? STATUS_SUCCESS : STATUS_UNSUCCESSFUL;
}

VOID ProcessSetMonitorItem(PIRP irp, PIO_STACK_LOCATION irpStack)
{
	// get data from buffer
	PHOOK_SET_MONITOR_ITEM_REQUEST pSetMonitorItemRequest = (PHOOK_SET_MONITOR_ITEM_REQUEST)irp->AssociatedIrp.SystemBuffer;
	PMONITOR_ITEM pMonitorItem = ExAllocatePool(NonPagedPool, sizeof(MONITOR_ITEM));
	pMonitorItem->ServiceTableBase = pSetMonitorItemRequest->ServiceTableBase;
	pMonitorItem->FunctionIndex = pSetMonitorItemRequest->FunctionIndex;
	pMonitorItem->ReloadedFunctionAddress = pSetMonitorItemRequest->ReloadedFunctionAddress;
	
	KIRQL oldIrq;
	KeAcquireSpinLock(&g_HookMonitorListSpinLock, &oldIrq);

	// insertion
	RemoveMonitorItem(pMonitorItem->FunctionIndex);
	if (pSetMonitorItemRequest->IsMonitor)
	{
		InsertTailList(&g_HookMonitorList, &pMonitorItem->ListEntry);
	}

	// response
	PLONG buffer = (PLONG)irp->AssociatedIrp.SystemBuffer;
	LONG count = 0;
	PLIST_ENTRY pList = g_HookMonitorList.Flink;
	while (pList != &g_HookMonitorList)
	{
		PMONITOR_ITEM pMonitorItem = CONTAINING_RECORD(pList, MONITOR_ITEM, ListEntry);
		*(buffer + count) = pMonitorItem->FunctionIndex;

		// next
		pList = pList->Flink;
		count++;
	}
	*(buffer + count) = -1; // last one

	KeReleaseSpinLock(&g_HookMonitorListSpinLock, oldIrq); 

	irp->IoStatus.Information = irpStack->Parameters.DeviceIoControl.OutputBufferLength;
	irp->IoStatus.Status = STATUS_SUCCESS;
}

VOID ProcessGetMonitorInfo(PIRP irp, PIO_STACK_LOCATION irpStack)
{
	BOOLEAN noMore = FALSE;
	ULONG outputBufferLength = irpStack->Parameters.DeviceIoControl.OutputBufferLength;
	PMONITOR_INFO_ITEM_DATA buffer = (PMONITOR_INFO_ITEM_DATA)irp->AssociatedIrp.SystemBuffer;
	KIRQL oldIrq;
	KeAcquireSpinLock(&g_HookMonitorInfoListSpinLock, &oldIrq);
	for (ULONG i = 0; i < outputBufferLength / sizeof(MONITOR_INFO_ITEM_DATA); i++)
	{
		PMONITOR_INFO_ITEM_DATA data = (buffer + i);
		if (!IsListEmpty(&g_HookMonitorInfoList))
		{
			PLIST_ENTRY pListEntry = RemoveHeadList(&g_HookMonitorInfoList);
			PMONITOR_INFO_ITEM pMonitorInfoItem = CONTAINING_RECORD(pListEntry, MONITOR_INFO_ITEM, ListEntry);
			*data = pMonitorInfoItem->Data; // copy data to buffer
			ExFreeToNPagedLookasideList(&g_HookMonitorInfoListLookaside, pMonitorInfoItem);
		}
		else
		{
			data->FunctionIndex = -1;
			noMore = TRUE;
			break;
		}
	}
	KeReleaseSpinLock(&g_HookMonitorInfoListSpinLock, oldIrq);

	if (noMore)
	{
		// reset event
		KeClearEvent(g_HookMonitorInfoListEvent);
	}

	irp->IoStatus.Information = outputBufferLength;
	irp->IoStatus.Status = STATUS_SUCCESS;
}

VOID EmptyHookLists()
{
	FreeList(&g_HookMonitorList, &g_HookMonitorListSpinLock, NULL, NULL);
	FreeList(&g_HookMonitorInfoList, &g_HookMonitorInfoListSpinLock, &g_HookMonitorInfoListLookaside, NULL);
}

VOID HookInternal(PVOID *pParams)
{
	// params
	ULONG addressOfHook = *(PULONG)pParams;
	ULONG addressOfWrite = *(PULONG)(pParams + 1);
	BOOLEAN isInstall = *(PBOOLEAN)(pParams + 2);

	if (isInstall)
	{
		// hook
		UCHAR jmpCodes[5];
		jmpCodes[0] = 0xe9;
		*(PULONG)&jmpCodes[1] = (ULONG)HookEntry - 5 - addressOfHook;
		g_HookReturnAddress = addressOfHook + 5;
		RtlCopyMemory((PVOID)addressOfWrite, jmpCodes, 5);
	}
	else
	{
		// unhook
		UCHAR orgCodes[5] = { 0x2b, 0xe1, 0xc1, 0xe9, 0x2 };
		RtlCopyMemory((PVOID)addressOfWrite, orgCodes, 5);
	}
}

__declspec(naked) VOID HookEntry()
{
	__asm
	{
		pushad
		pushfd

		push esi // paramAddress
		push ecx // paramByteCount
		push edx // orgFuncAddress
		push eax // funcIndex
		push edi // serviceTableBase
		call HookFilter
		mov[esp + 0x18], eax // change function address edx

		popfd
		popad

		sub esp, ecx // old code 1
		shr ecx, 2 // old code 2
		jmp g_HookReturnAddress 
	}
}

ULONG HookFilter(ULONG serviceTableBase, ULONG funcIndex, ULONG orgFuncAddress, ULONG paramByteCount, PULONG paramAddress)
{
	ULONG funcAddress = orgFuncAddress;
	ULONG paramCount = paramByteCount >> 2;
	KIRQL oldIrq;
	KeAcquireSpinLock(&g_HookMonitorListSpinLock, &oldIrq);

	PLIST_ENTRY pList = g_HookMonitorList.Flink;
	while (pList != &g_HookMonitorList)
	{
		PMONITOR_ITEM pMonitorItem = CONTAINING_RECORD(pList, MONITOR_ITEM, ListEntry);
		if (pMonitorItem->ServiceTableBase == serviceTableBase && pMonitorItem->FunctionIndex == funcIndex)
		{
			KdPrint(("service table 0x%x, index 0x%x, function 0x%x, param count %d at 0x%x", serviceTableBase, funcIndex, orgFuncAddress, paramCount, paramAddress));

			// replace function
			if (pMonitorItem->ReloadedFunctionAddress != 0 && MmIsAddressValid((PVOID)pMonitorItem->ReloadedFunctionAddress))
			{
				funcAddress = pMonitorItem->ReloadedFunctionAddress;
				KdPrint(("replaced function 0x%x with 0x%x", orgFuncAddress, funcAddress));
			}

			// collect info
			PMONITOR_INFO_ITEM pMonitorInfoItem = ExAllocateFromNPagedLookasideList(&g_HookMonitorInfoListLookaside);
			pMonitorInfoItem->Data.ProcessId = (ULONG)PsGetCurrentProcessId(); // caller pid
			pMonitorInfoItem->Data.FunctionIndex = funcIndex;
			pMonitorInfoItem->Data.FunctionAddress = funcAddress;
			pMonitorInfoItem->Data.ParameterCount = paramCount;
			for (ULONG i = 0; i < paramCount; i++)
			{
				pMonitorInfoItem->Data.Parameters[i] = *(paramAddress + i);
			}

			// insert monitor info
			KIRQL oldIrq;
			KeAcquireSpinLock(&g_HookMonitorInfoListSpinLock, &oldIrq);
			InsertTailList(&g_HookMonitorInfoList, &pMonitorInfoItem->ListEntry);
			KeReleaseSpinLock(&g_HookMonitorInfoListSpinLock, oldIrq);

			// signal event
			KeSetEvent(g_HookMonitorInfoListEvent, 0, FALSE);

			break;
		}
		pList = pList->Flink;
	}

	KeReleaseSpinLock(&g_HookMonitorListSpinLock, oldIrq);
	return funcAddress;
}

BOOLEAN RemoveMonitorItem(ULONG functionIndex)
{
	PLIST_ENTRY pList = g_HookMonitorList.Flink;
	while (pList != &g_HookMonitorList)
	{
		PMONITOR_ITEM pMonitorItem = CONTAINING_RECORD(pList, MONITOR_ITEM, ListEntry);
		if (pMonitorItem->FunctionIndex == functionIndex)
		{
			KdPrint(("removing monitor function index 0x%x", functionIndex));
			return RemoveEntryList(pList);
		}
		pList = pList->Flink;
	}
	return FALSE;
}

